import os
import ujson
import aicube
import time
from libs.PipeLine import ScopedTiming  # 用于性能计时的上下文管理器
from libs.Utils import *  # 导入自定义工具函数
from media.sensor import *  # 摄像头传感器相关功能
from media.display import *  # 显示设备相关功能
from media.media import *  # 媒体处理基础功能
import nncase_runtime as nn  # K210 AI运行时库
import ulab.numpy as np  # 嵌入式优化的numpy库
import image  # 图像处理模块
import gc  # 垃圾回收模块

# ================== 系统配置 ==================
# 显示模式配置（LCD或HDMI）
display_mode = "lcd"
if display_mode == "lcd":
    # LCD分辨率（800x480），按16字节对齐
    DISPLAY_WIDTH = ALIGN_UP(800, 16)
    DISPLAY_HEIGHT = 480
else:
    # HDMI分辨率（1920x1080），按16字节对齐
    DISPLAY_WIDTH = ALIGN_UP(1920, 16)
    DISPLAY_HEIGHT = 1080

# AI处理使用的图像分辨率（1280x720），按16字节对齐
OUT_RGB888P_WIDTH = ALIGN_UP(1280, 16)
OUT_RGB888P_HEIGH = 720

# 项目根路径和配置文件路径
root_path = "/sdcard/mp_deployment_source/"
config_path = root_path + "deploy_config.json"
debug_mode = 1  # 调试模式开关（1为开启）

# ================== 污渍类型配置 ==================
# 污渍类型及其清洁优先级（数值越小优先级越高）
STAIN_PRIORITY = {
    "Droppings": 1,   # 最高优先级 - 鸟类排泄物
    "Leaves": 2,      # 高优先级 - 落叶
    "Branch": 3,      # 中优先级 - 树枝
    "Paper": 4        # 低优先级 - 纸张
}

# 清洁建议映射
CLEANING_ADVICE = {
    "Droppings": "立即清洁！具有腐蚀性",
    "Leaves": "尽快清理，避免堆积",
    "Branch": "中等优先级，可安排清理",
    "Paper": "低优先级，最后处理"
}

# 优先级颜色映射（用于UI显示）
PRIORITY_COLORS = {
    1: (255, 0, 0),      # 红色 - 最高优先级
    2: (255, 165, 0),    # 橙色
    3: (255, 255, 0),    # 黄色
    4: (0, 255, 0)       # 绿色 - 最低优先级
}

def two_side_pad_param(input_size, output_size):
    """
    计算保持长宽比的双边填充参数

    Args:
        input_size: 原始图像尺寸 (宽, 高)
        output_size: 目标图像尺寸 (宽, 高)

    Returns:
        top, bottom, left, right: 各方向填充像素数
        ratio: 缩放比例
    """
    # 计算宽高缩放比例
    ratio_w = output_size[0] / input_size[0]
    ratio_h = output_size[1] / input_size[1]
    ratio = min(ratio_w, ratio_h)  # 取最小比例保证图像完整

    # 计算缩放后的新尺寸
    new_w = int(ratio * input_size[0])
    new_h = int(ratio * input_size[1])

    # 计算需要填充的像素
    dw = (output_size[0] - new_w) / 2
    dh = (output_size[1] - new_h) / 2

    # 四边填充像素计算（取整）
    top = int(round(dh - 0.1))
    bottom = int(round(dh + 0.1))
    left = int(round(dw - 0.1))
    right = int(round(dw - 0.1))

    return top, bottom, left, right, ratio

def read_deploy_config(config_path):
    """
    读取部署配置文件

    Args:
        config_path: 配置文件路径

    Returns:
        dict: 解析后的配置字典
    """
    with open(config_path, 'r') as f:
        try:
            # 使用ujson快速解析JSON配置
            return ujson.load(f)
        except Exception as e:
            print(f"配置文件解析错误: {e}")
            return {}  # 返回空字典避免后续错误

def get_priority_info(det_boxes, labels):
    """
    评估污渍的清洁优先级和统计数量

    Args:
        det_boxes: 检测到的污渍框列表
        labels: 污渍类型标签列表

    Returns:
        highest_priority: 当前最高优先级值
        highest_priority_stain: 最高优先级污渍类型
        priority_counts: 各优先级污渍数量统计字典
    """
    # 初始化优先级计数器
    priority_counts = {1: 0, 2: 0, 3: 0, 4: 0}
    highest_priority = 5  # 初始化为大于最高优先级的值
    highest_priority_stain = "无污渍"

    # 遍历所有检测框
    for box in det_boxes:
        # 获取污渍类型
        stain_type = labels[box[0]]

        # 获取该污渍类型的优先级（默认为最低优先级4）
        priority = STAIN_PRIORITY.get(stain_type, 4)

        # 更新最高优先级
        if priority < highest_priority:
            highest_priority = priority
            highest_priority_stain = stain_type

        # 统计各优先级污渍数量
        if priority in priority_counts:
            priority_counts[priority] += 1

    return highest_priority, highest_priority_stain, priority_counts

def detection():
    """目标检测系统主函数"""
    print("目标检测系统启动...")

    # ================== 初始化阶段 ==================
    # 1. 加载部署配置
    deploy_conf = read_deploy_config(config_path)
    if not deploy_conf:
        print("错误：无法加载部署配置，系统退出")
        return -1

    # 从配置中获取参数
    kmodel_name = deploy_conf["kmodel_path"]  # 模型文件名
    labels = deploy_conf["categories"]        # 类别标签列表
    confidence_threshold = deploy_conf["confidence_threshold"]  # 置信度阈值
    nms_threshold = deploy_conf["nms_threshold"]  # NMS阈值
    img_size = deploy_conf["img_size"]        # 模型输入尺寸
    num_classes = deploy_conf["num_classes"]   # 类别数量
    color_four = get_colors(num_classes)       # 为每个类别生成颜色
    nms_option = deploy_conf["nms_option"]     # NMS算法选项
    model_type = deploy_conf["model_type"]     # 模型类型

    # 如果是基于锚点的检测模型，加载锚点配置
    if model_type == "AnchorBaseDet":
        anchors = deploy_conf["anchors"][0] + deploy_conf["anchors"][1] + deploy_conf["anchors"][2]
    else:
        anchors = []  # 非锚点模型不需要锚点

    # 模型输入尺寸和原始图像尺寸
    kmodel_frame_size = img_size
    frame_size = [OUT_RGB888P_WIDTH, OUT_RGB888P_HEIGH]
    strides = [8, 16, 32]  # 特征图步长（用于后处理）

    # 2. 计算预处理填充参数（保持长宽比）
    top, bottom, left, right, ratio = two_side_pad_param(frame_size, kmodel_frame_size)

    # 3. 初始化AI推理引擎（KPU）
    kpu = nn.kpu()
    kpu.load_kmodel(root_path + kmodel_name)  # 加载kmodel模型

    # 4. 初始化AI2D预处理引擎
    ai2d = nn.ai2d()
    # 设置输入输出格式（NCHW）和数据类型
    ai2d.set_dtype(nn.ai2d_format.NCHW_FMT, nn.ai2d_format.NCHW_FMT, np.uint8, np.uint8)
    # 配置填充参数（保持比例，填充灰色[114,114,114]）
    ai2d.set_pad_param(True, [0, 0, 0, 0, top, bottom, left, right], 0, [114, 114, 114])
    # 设置缩放参数（双线性插值）
    ai2d.set_resize_param(True, nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
    # 构建AI2D计算图
    ai2d_builder = ai2d.build(
        [1, 3, OUT_RGB888P_HEIGH, OUT_RGB888P_WIDTH],  # 输入张量形状
        [1, 3, kmodel_frame_size[1], kmodel_frame_size[0]]  # 输出张量形状
    )

    # 5. 初始化摄像头传感器
    sensor = Sensor()
    sensor.reset()  # 重置传感器
    sensor.set_hmirror(False)  # 关闭水平镜像
    sensor.set_vflip(False)    # 关闭垂直翻转

    # 配置通道0（用于显示）：YUV420格式
    sensor.set_framesize(width=DISPLAY_WIDTH, height=DISPLAY_HEIGHT)
    sensor.set_pixformat(PIXEL_FORMAT_YUV_SEMIPLANAR_420)

    # 配置通道2（用于AI处理）：RGB888格式
    sensor.set_framesize(width=OUT_RGB888P_WIDTH, height=OUT_RGB888P_HEIGH, chn=CAM_CHN_ID_2)
    sensor.set_pixformat(PIXEL_FORMAT_RGB_888_PLANAR, chn=CAM_CHN_ID_2)

    # 6. 配置显示系统
    # 绑定摄像头通道0到显示层
    sensor_bind_info = sensor.bind_info(x=0, y=0, chn=CAM_CHN_ID_0)
    Display.bind_layer(**sensor_bind_info, layer=Display.LAYER_VIDEO1)

    # 初始化显示设备
    if display_mode == "lcd":
        Display.init(Display.ST7701, to_ide=True)  # LCD显示屏
    else:
        Display.init(Display.LT9611, to_ide=True)  # HDMI显示屏

    # 创建OSD图像层（用于绘制检测结果）
    osd_img = image.Image(DISPLAY_WIDTH, DISPLAY_HEIGHT, image.ARGB8888)

    # 7. 初始化媒体管理器
    MediaManager.init()

    # 8. 启动摄像头
    sensor.run()

    # 9. 初始化AI2D输出张量（预分配内存）
    data = np.ones((1, 3, kmodel_frame_size[1], kmodel_frame_size[0]), dtype=np.uint8)
    ai2d_output_tensor = nn.from_numpy(data)

    # ================== 主循环 ==================
    print("进入主循环...")
    while True:
        # 使用性能计时器（仅在调试模式下启用）
        with ScopedTiming("total", debug_mode > 0):
            # 从通道2捕获RGB图像
            rgb888p_img = sensor.snapshot(chn=CAM_CHN_ID_2)

            # 检查图像格式是否正确
            if rgb888p_img.format() == image.RGBP888:
                # 将图像转为numpy数组
                ai2d_input = rgb888p_img.to_numpy_ref()
                ai2d_input_tensor = nn.from_numpy(ai2d_input)

                # 执行AI2D预处理（缩放+填充）
                ai2d_builder.run(ai2d_input_tensor, ai2d_output_tensor)

                # 设置模型输入
                kpu.set_input_tensor(0, ai2d_output_tensor)

                # 模型推理
                kpu.run()

                # 获取输出结果
                results = []
                for i in range(kpu.outputs_size()):
                    out_data = kpu.get_output_tensor(i)
                    # 展平输出结果（1D数组）
                    result = out_data.to_numpy().flatten()
                    del out_data  # 及时释放张量内存
                    results.append(result)

                # 后处理：生成检测框
                det_boxes = aicube.anchorbasedet_post_process(
                    results[0], results[1], results[2],
                    kmodel_frame_size, frame_size, strides,
                    num_classes, confidence_threshold,
                    nms_threshold, anchors, nms_option
                )

                # 污渍优先级评估
                highest_priority = 5  # 初始值（大于所有优先级）
                highest_stain = "无污渍"
                priority_counts = {1: 0, 2: 0, 3: 0, 4: 0}  # 各优先级计数器

                # 如果有检测结果，进行优先级评估
                if det_boxes:
                    highest_priority, highest_stain, priority_counts = get_priority_info(det_boxes, labels)

                # 绘制UI
                osd_img.clear()  # 清空上一帧结果

                # ========== 顶部状态栏绘制 ==========
                status_bar_height = 60
                # 绘制半透明黑色背景状态栏
                osd_img.draw_rectangle(0, 0, DISPLAY_WIDTH, status_bar_height,
                                      color=(0, 0, 0, 180), fill=True)

                # 显示清洁优先级
                priority_color = PRIORITY_COLORS.get(highest_priority, (255, 255, 255))
                priority_text = f"清洁优先级: {highest_stain} (等级{highest_priority})"
                osd_img.draw_string_advanced(10, 10, 32, priority_text, color=priority_color)

                # 显示清洁建议（仅在检测到污渍时显示）
                if highest_stain != "无污渍":
                    advice = CLEANING_ADVICE.get(highest_stain, "建议清理")
                    advice_text = f"建议: {advice}"
                    osd_img.draw_string_advanced(10, 40, 24, advice_text, color=(255, 255, 255))

                # ========== 污渍统计面板 ==========
                stats_x = DISPLAY_WIDTH - 150  # 右侧起始位置
                stats_y = 10
                # 绘制统计标题
                osd_img.draw_string_advanced(stats_x, stats_y, 24, "污渍统计:", color=(200, 200, 200))
                stats_y += 30

                # 显示各优先级污渍数量
                for prio in sorted(priority_counts.keys()):
                    count = priority_counts[prio]
                    if count > 0:
                        color = PRIORITY_COLORS.get(prio, (255, 255, 255))
                        # 获取污渍类型名称（取首字母）
                        stain_name = [k for k, v in STAIN_PRIORITY.items() if v == prio][0][0]
                        stat_text = f"P{prio}: {stain_name}x{count}"
                        osd_img.draw_string_advanced(stats_x, stats_y, 24, stat_text, color=color)
                        stats_y += 30

                # ========== 绘制检测框和标签 ==========
                if det_boxes:
                    for box in det_boxes:
                        # 解析检测框坐标 [x1, y1, x2, y2]
                        x1, y1, x2, y2 = box[2], box[3], box[4], box[5]

                        # 获取污渍信息
                        stain_type = labels[box[0]]
                        confidence = box[1]
                        priority = STAIN_PRIORITY.get(stain_type, 4)
                        color = PRIORITY_COLORS.get(priority, (255, 255, 255))

                        # 坐标映射到显示分辨率（考虑状态栏偏移）
                        x = int(x1 * DISPLAY_WIDTH / OUT_RGB888P_WIDTH)
                        y = int(y1 * DISPLAY_HEIGHT / OUT_RGB888P_HEIGH) + status_bar_height
                        w = int((x2 - x1) * DISPLAY_WIDTH / OUT_RGB888P_WIDTH)
                        h = int((y2 - y1) * DISPLAY_HEIGHT / OUT_RGB888P_HEIGH)

                        # 绘制检测框（使用优先级颜色）
                        osd_img.draw_rectangle(x, y, w, h, color=color, thickness=3)

                        # 绘制标签（污渍类型 + 置信度）
                        label_text = f"{stain_type} {confidence:.2f}"
                        osd_img.draw_string_advanced(x, y - 30, 28, label_text, color=color)

                # 显示结果到OSD层
                Display.show_image(osd_img, 0, 0, Display.LAYER_OSD3)

                # ========== 控制台输出 ==========
                print("\n" + "=" * 50)
                print(f"检测到污渍数量: {len(det_boxes)}")

                if det_boxes:
                    print(f"最高优先级污渍: {highest_stain} (等级{highest_priority})")
                    print("污渍分布统计:")
                    for prio in sorted(priority_counts.keys()):
                        count = priority_counts[prio]
                        if count > 0:
                            stain_name = [k for k, v in STAIN_PRIORITY.items() if v == prio][0]
                            print(f"  等级{prio}: {stain_name} - {count}个")
                    print(f"清洁建议: {CLEANING_ADVICE.get(highest_stain, '常规清理')}")
                else:
                    print("未检测到污渍，区域清洁")

                # 手动触发垃圾回收
                gc.collect()

            # 释放当前帧图像引用
            rgb888p_img = None

    # ================== 清理阶段 ==================
    # 注意：正常情况下不会执行到这里（无限循环）
    print("系统正在关闭...")
    sensor.stop()           # 停止摄像头
    Display.deinit()        # 释放显示资源
    MediaManager.deinit()   # 释放媒体资源

    # 释放内存
    del ai2d_output_tensor
    gc.collect()
    time.sleep(1)
    nn.shrink_memory_pool()  # 释放KPU内存池

    print("系统已关闭")
    return 0

if __name__ == "__main__":
    # 启动检测系统
    detection()
